from pandas import read_csv
from pandas import DataFrame
from pandas import concat
from matplotlib import pyplot
from sklearn.preprocessing import MinMaxScaler
from keras.models import Sequential
from keras.layers import LSTM
from keras.layers import Dense
from keras.layers import Dropout


def series_to_supervised(data, n_in=1, n_out=1, dropnan=True):
    n_vars = 1 if type(data) is list else data.shape[1]
    df = DataFrame(data)
    cols, names = list(), list()
    # input sequence (t-n, ... t-1)
    for i in range(n_in, 0, -1):
        cols.append(df.shift(i))
        names += [('var%d(t-%d)' % (j+1, i)) for j in range(n_vars)]
    # forecast sequence (t, t+1, ... t+n)
    for i in range(0, n_out):
        cols.append(df.shift(-i))
        if i == 0:
            names += [('var%d(t)' % (j+1)) for j in range(n_vars)]
        else:
            names += [('var%d(t+%d)' % (j+1, i)) for j in range(n_vars)]
    # put it all together
    agg = concat(cols, axis=1)
    agg.columns = names
    # drop rows with NaN values
    if dropnan:
        agg.dropna(inplace=True)
    return agg


if __name__ == "__main__":
    #数据的读取和处理
    dataset = read_csv('bull_data2_c.csv', index_col=0)
    dataset.index.name = 'data'
    values = dataset.values
    values = values.astype('float32')
    scaler = MinMaxScaler(feature_range=(0, 1)) #归一化
    scaled = scaler.fit_transform(values)
    reframed = series_to_supervised(scaled, 1, 1) #转变为监督学习
    values_ss = reframed.values
    train_data = 2500 #训练集数目
    #划分训练集和测试集
    train = values_ss[:train_data, :]
    test = values_ss[train_data:, :]
    train_X = train[:, :7]
    train_Y = train[:, 7:]
    test_X = test[:, :7]
    test_Y = test[:, 7:]
    #升维成三维张量
    train_X = train_X.reshape((train_X.shape[0], 1, train_X.shape[1]))
    test_X = test_X.reshape((test_X.shape[0], 1, test_X.shape[1]))
    #模型搭建
    model = Sequential()
    model.add(LSTM(50, input_shape=(train_X.shape[1], train_X.shape[2])))
    model.add(Dense(7))
    model.add(Dropout(0.25))
    model.compile(loss='mean_squared_error', optimizer='adam')
    model.fit(train_X, train_Y, epochs=50, batch_size=1, validation_data=(test_X, test_Y), verbose=1, shuffle=False)
    model.save("my_model.h5")
    #利用模型进行预测
    yhat = model.predict(test_X)
    yhat_change = scaler.inverse_transform(yhat)
    test_X = test_X.reshape((test_X.shape[0], test_X.shape[2]))
    test_Y = scaler.inverse_transform(test_Y)
    for i in range(test_X.shape[0]):
        pyplot.plot(yhat_change[i])
        pyplot.plot(test_Y[i])
    pyplot.show()
